MODULE CMF_CALC_FLDSTG_MOD
!==========================================================
!* PURPOSE: calculate river and floodplain staging
!
! (C) D.Yamazaki & E. Dutra  (U-Tokyo/FCUL)  Aug 2019
!
! Licensed under the Apache License, Version 2.0 (the "License");
!   You may not use this file except in compliance with the License.
!   You may obtain a copy of the License at: http://www.apache.org/licenses/LICENSE-2.0
!
! Unless required by applicable law or agreed to in writing, software distributed under the License is 
!  distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
! See the License for the specific language governing permissions and limitations under the License.
!==========================================================
CONTAINS
!####################################################################
! -- CMF_CALC_FLDSTG_DEF  !! default flood stage calculation
! -- CMF_OPT_FLDSTG_ES    !! optimized code for vector processor (such as Earth Simulator), activated using LSTG_ES=.TRUE. option).
! -- CMP_OPT_FLDSTG_SIMD  !! SIMD implementation primary for OpenMP (not fast enough)
!####################################################################
SUBROUTINE CMF_CALC_FLDSTG_DEF
USE PARKIND1,           ONLY: JPIM, JPRB, JPRD
USE YOS_CMF_INPUT,      ONLY: NLFP
USE YOS_CMF_MAP,        ONLY: NSEQALL, D2GRAREA, D2RIVLEN, D2RIVWTH, D2RIVELV
USE YOS_CMF_MAP,        ONLY: D2RIVSTOMAX, D2FLDSTOMAX, D2FLDGRD, DFRCINC
USE YOS_CMF_PROG,       ONLY: P2RIVSTO, P2FLDSTO
USE YOS_CMF_DIAG,       ONLY: D2RIVDPH, D2FLDDPH, D2FLDFRC, D2FLDARE, D2SFCELV, D2STORGE
USE YOS_CMF_DIAG,       ONLY: P0GLBSTOPRE2, P0GLBSTONEW2, P0GLBRIVSTO, P0GLBFLDSTO, P0GLBFLDARE
IMPLICIT NONE

!*** LOCAL
INTEGER(KIND=JPIM),SAVE    :: ISEQ, I
REAL(KIND=JPRD),SAVE       :: PSTOALL
REAL(KIND=JPRB),SAVE       :: DSTOALL,DSTO_add,DSTO_fil,DWTH_add,DWTH_fil,DDPH_fil,DWTH_inc
!$OMP THREADPRIVATE        (I,DSTOALL,DSTO_add,DSTO_fil,DWTH_add,DWTH_fil,DDPH_fil,DWTH_inc,PSTOALL)
!================================================
P0GLBSTOPRE2=0._JPRD
P0GLBSTONEW2=0._JPRD
P0GLBRIVSTO =0._JPRD
P0GLBFLDSTO =0._JPRD
P0GLBFLDARE =0._JPRD

! Estimate water depth and flood extent from water storage
!   Solution for Equations (1) and (2) in [Yamazaki et al. 2011 WRR].

!$OMP PARALLEL DO REDUCTION(+:P0GLBSTOPRE2)
DO ISEQ=1, NSEQALL
!
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)
  P0GLBSTOPRE2     = P0GLBSTOPRE2 + PSTOALL

  IF( PSTOALL > D2RIVSTOMAX(ISEQ,1) )THEN
    I=1
    DSTO_fil = D2RIVSTOMAX(ISEQ,1)
    DWTH_fil = D2RIVWTH(ISEQ,1)
    DDPH_fil = 0._JPRB
    DWTH_inc = D2GRAREA(ISEQ,1) / D2RIVLEN(ISEQ,1) * DFRCINC
    DO WHILE( PSTOALL > D2FLDSTOMAX(ISEQ,1,I) .AND. I<=NLFP)
      DSTO_fil = D2FLDSTOMAX(ISEQ,1,I)
      DWTH_fil = DWTH_fil + DWTH_inc
      DDPH_fil = DDPH_fil + D2FLDGRD(ISEQ,1,I) * DWTH_inc
      I=I+1
      IF( I>NLFP ) EXIT
    ENDDO
    IF( I>NLFP )THEN
      DSTO_add = DSTOALL - DSTO_fil
      DWTH_add = 0._JPRB
      D2FLDDPH(ISEQ,1) = DDPH_fil + DSTO_add / DWTH_fil / D2RIVLEN(ISEQ,1)
    ELSE
      DSTO_add =  DSTOALL - DSTO_fil
      DWTH_add = -DWTH_fil + &
&      ( DWTH_fil**2._JPRB + 2._JPRB * DSTO_add / D2RIVLEN(ISEQ,1) / D2FLDGRD(ISEQ,1,I) )**0.5_JPRB
      D2FLDDPH(ISEQ,1) = DDPH_fil + D2FLDGRD(ISEQ,1,I) * DWTH_add
    ENDIF
    P2RIVSTO(ISEQ,1) = D2RIVSTOMAX(ISEQ,1) + D2RIVLEN(ISEQ,1) * D2RIVWTH(ISEQ,1) * D2FLDDPH(ISEQ,1)
    P2RIVSTO(ISEQ,1) = MIN(P2RIVSTO(ISEQ,1),PSTOALL)

    D2RIVDPH(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1),JPRB) / D2RIVLEN(ISEQ,1) / D2RIVWTH(ISEQ,1)
!
    P2FLDSTO(ISEQ,1) = PSTOALL - P2RIVSTO(ISEQ,1)
    P2FLDSTO(ISEQ,1) = MAX( P2FLDSTO(ISEQ,1), 0._JPRD )
    D2FLDFRC(ISEQ,1) = (-D2RIVWTH(ISEQ,1) + DWTH_fil + DWTH_add ) / (DWTH_inc*NLFP)
    D2FLDFRC(ISEQ,1) = MAX( D2FLDFRC(ISEQ,1),0._JPRB)
    D2FLDFRC(ISEQ,1) = MIN( D2FLDFRC(ISEQ,1),1._JPRB)
    D2FLDARE(ISEQ,1) = D2GRAREA(ISEQ,1)*D2FLDFRC(ISEQ,1)
  ELSE
    P2RIVSTO(ISEQ,1) = PSTOALL
    D2RIVDPH(ISEQ,1) = DSTOALL / D2RIVLEN(ISEQ,1) / D2RIVWTH(ISEQ,1)
    D2RIVDPH(ISEQ,1) = MAX( D2RIVDPH(ISEQ,1), 0._JPRB )
    P2FLDSTO(ISEQ,1) = 0._JPRD
    D2FLDDPH(ISEQ,1) = 0._JPRB
    D2FLDFRC(ISEQ,1) = 0._JPRB
    D2FLDARE(ISEQ,1) = 0._JPRB
  ENDIF
ENDDO
!$OMP END PARALLEL DO

!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTONEW2,P0GLBRIVSTO,P0GLBFLDSTO,P0GLBFLDARE)
DO ISEQ=1, NSEQALL
  D2SFCELV(ISEQ,1) = D2RIVELV(ISEQ,1) + D2RIVDPH(ISEQ,1)
  D2STORGE(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1)+P2FLDSTO(ISEQ,1),KIND=JPRB)
  P0GLBSTONEW2     = P0GLBSTONEW2 + P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBRIVSTO      = P0GLBRIVSTO  + P2RIVSTO(ISEQ,1)
  P0GLBFLDSTO      = P0GLBFLDSTO  + P2FLDSTO(ISEQ,1)
  P0GLBFLDARE      = P0GLBFLDARE  + D2FLDARE(ISEQ,1)
ENDDO
!$OMP END PARALLEL DO SIMD

END SUBROUTINE CMF_CALC_FLDSTG_DEF
!####################################################################
!
!
!
!
!####################################################################
SUBROUTINE CMF_OPT_FLDSTG_SIMD
USE PARKIND1,           ONLY: JPIM, JPRB, JPRD
USE YOS_CMF_INPUT,      ONLY: NLFP
USE YOS_CMF_MAP,        ONLY: NSEQALL, D2GRAREA, D2RIVLEN, D2RIVWTH, D2RIVELV
USE YOS_CMF_MAP,        ONLY: D2RIVSTOMAX, D2FLDSTOMAX, D2FLDHGT, D2FLDGRD, DFRCINC
USE YOS_CMF_PROG,       ONLY: P2RIVSTO, P2FLDSTO
USE YOS_CMF_DIAG,       ONLY: D2RIVDPH, D2FLDDPH, D2FLDFRC, D2FLDARE, D2SFCELV, D2STORGE
USE YOS_CMF_DIAG,       ONLY: P0GLBSTOPRE2, P0GLBSTONEW2, P0GLBRIVSTO, P0GLBFLDSTO, P0GLBFLDARE
IMPLICIT NONE

!*** LOCAL
INTEGER(KIND=JPIM)         :: ISTG(NSEQALL)
REAL(KIND=JPRB)            :: DGRD(NSEQALL)
REAL(KIND=JPRB)            :: DSTO_fil(NSEQALL)
REAL(KIND=JPRB)            :: DDPH_fil(NSEQALL)

INTEGER(KIND=JPIM),SAVE    :: ISEQ, I, ILFP
LOGICAL,SAVE               :: Mask
REAL(KIND=JPRD),SAVE       :: PSTOALL
REAL(KIND=JPRB),SAVE       :: DSTOALL, DWTH_fil, DSTO_add, DWTH_add, Dmod, Dtmp
!================================================
P0GLBSTOPRE2=0._JPRD
P0GLBSTONEW2=0._JPRD
P0GLBRIVSTO =0._JPRD
P0GLBFLDSTO =0._JPRD
P0GLBFLDARE =0._JPRD

! Estimate water depth and flood extent from water storage
!   Solution for Equations (1) and (2) in [Yamazaki et al. 2011 WRR].

!$OMP PARALLEL DO PRIVATE(I,ILFP,PSTOALL) REDUCTION(+:P0GLBSTOPRE2)
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBSTOPRE2 = P0GLBSTOPRE2 + PSTOALL

  I=0 !! all water in channel
  DSTO_fil(ISEQ)=0._JPRB
  DDPH_fil(ISEQ)=0._JPRB
  IF (PSTOALL > D2RIVSTOMAX(ISEQ,1)) THEN
    I=1
    DSTO_fil(ISEQ)=D2RIVSTOMAX(ISEQ,1)
    DDPH_fil(ISEQ)=0._JPRB
    DO ILFP=1, NLFP
      IF (PSTOALL <= D2FLDSTOMAX(ISEQ,1,ILFP)) EXIT
      I=ILFP+1 !! water in river and ith floodplain layer
      DSTO_fil(ISEQ)=D2FLDSTOMAX(ISEQ,1,I-1)
      DDPH_fil(ISEQ)=D2FLDHGT(ISEQ,1,I-1)
    ENDDO
  ENDIF
  IF( I==0 .or. I==NLFP+1 )THEN
    DGRD(ISEQ)=1.E10_JPRB
  ELSE
    DGRD(ISEQ)=D2FLDGRD(ISEQ,1,I)
  ENDIF
ENDDO
!$OMP END PARALLEL DO 

!$OMP PARALLEL DO SIMD PRIVATE(PSTOALL,DSTOALL,DWTH_fil,DSTO_add,DWTH_add,Dtmp,Dmod,Mask)
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)

  DWTH_fil = D2GRAREA(ISEQ,1) / D2RIVLEN(ISEQ,1) * DFRCINC
  DWTH_fil=DWTH_fil*ISTG(ISEQ)

  !! calculate storage, width, depth increment for regular case ( 0<I<=NLFP )
  DSTO_add =  DSTOALL - DSTO_fil(ISEQ)
  DWTH_add = -DWTH_fil + ( DWTH_fil**2. + 2._JPRB * DSTO_add / D2RIVLEN(ISEQ,1) / DGRD(ISEQ) )**0.5_JPRB
  D2FLDDPH(ISEQ,1) = DDPH_fil(ISEQ) +DGRD(ISEQ) * DWTH_add

  !! Case when all catchment is flooded
  Mask = ( ISTG(ISEQ)<=NLFP ) 
   DWTH_add=merge(DWTH_add,0._JPRB,Mask)
   Dtmp=D2FLDDPH(ISEQ,1)
   Dmod=DDPH_fil(ISEQ) + DSTO_add / DWTH_fil / D2RIVLEN(ISEQ,1)
   D2FLDDPH(ISEQ,1)=merge(Dtmp,Dmod,Mask)

  !! Case when all water stay in channel
  Mask = ( ISTG(ISEQ) >0 ) 
   Dmod=D2RIVWTH(ISEQ,1)
   DWTH_add=merge(DWTH_add,Dmod,Mask)
   Dtmp=D2FLDDPH(ISEQ,1)
   D2FLDDPH(ISEQ,1)=merge(Dtmp,0._JPRB,Mask)

  P2RIVSTO(ISEQ,1) = D2RIVSTOMAX(ISEQ,1) + D2RIVLEN(ISEQ,1) * D2RIVWTH(ISEQ,1) * D2FLDDPH(ISEQ,1)
  P2RIVSTO(ISEQ,1) = MIN(P2RIVSTO(ISEQ,1),PSTOALL)
  D2RIVDPH(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1),JPRB) / D2RIVLEN(ISEQ,1) / D2RIVWTH(ISEQ,1)
  D2RIVDPH(ISEQ,1) = MAX( D2RIVDPH(ISEQ,1), 0._JPRB )
!
  P2FLDSTO(ISEQ,1) = PSTOALL - P2RIVSTO(ISEQ,1)
  P2FLDSTO(ISEQ,1) = MAX( P2FLDSTO(ISEQ,1), 0._JPRD )

  D2FLDFRC(ISEQ,1) = (-D2RIVWTH(ISEQ,1)+DWTH_fil+DWTH_add)*D2RIVLEN(ISEQ,1) / D2GRAREA(ISEQ,1)
  D2FLDFRC(ISEQ,1) = MAX( D2FLDFRC(ISEQ,1),0._JPRB)
  D2FLDFRC(ISEQ,1) = MIN( D2FLDFRC(ISEQ,1),1._JPRB)
  D2FLDARE(ISEQ,1) = D2GRAREA(ISEQ,1)*D2FLDFRC(ISEQ,1)

  D2SFCELV(ISEQ,1) = D2RIVELV(ISEQ,1) + D2RIVDPH(ISEQ,1)
  D2STORGE(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1)+P2FLDSTO(ISEQ,1),KIND=JPRB)
ENDDO
!$OMP END PARALLEL DO SIMD

!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTONEW2,P0GLBRIVSTO,P0GLBFLDSTO,P0GLBFLDARE)
DO ISEQ=1, NSEQALL
  P0GLBSTONEW2     = P0GLBSTONEW2 + P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBRIVSTO      = P0GLBRIVSTO  + P2RIVSTO(ISEQ,1)
  P0GLBFLDSTO      = P0GLBFLDSTO  + P2FLDSTO(ISEQ,1)
  P0GLBFLDARE      = P0GLBFLDARE  + D2FLDARE(ISEQ,1)
ENDDO
!$OMP END PARALLEL DO SIMD

END SUBROUTINE CMF_OPT_FLDSTG_SIMD
!####################################################################
!
!
!
!
!
!####################################################################
SUBROUTINE CMF_OPT_FLDSTG_ES
! ==========
! Optional code for Earth Simulator (Vector Processor)
! Specify option: LSTG_ES=.TRUE.
! Faster computation on vector prosessor by avoiding IF-THEN function. (note this code will be slow on Scaler Processor)
! ==========
USE PARKIND1,           ONLY: JPIM, JPRB, JPRD
USE YOS_CMF_INPUT,      ONLY: NLFP
USE YOS_CMF_MAP,        ONLY: NSEQALL, D2GRAREA, D2RIVLEN, D2RIVWTH, D2RIVELV
USE YOS_CMF_MAP,        ONLY: D2RIVSTOMAX, D2FLDSTOMAX, D2FLDGRD, DFRCINC
USE YOS_CMF_PROG,       ONLY: P2RIVSTO, P2FLDSTO
USE YOS_CMF_DIAG,       ONLY: D2RIVDPH, D2FLDDPH, D2FLDFRC, D2FLDARE, D2SFCELV, D2STORGE
USE YOS_CMF_DIAG,       ONLY: P0GLBSTOPRE2, P0GLBSTONEW2, P0GLBRIVSTO, P0GLBFLDSTO, P0GLBFLDARE
IMPLICIT NONE

!*** LOCAL
REAL(KIND=JPRB)            :: D2STO_fil(NSEQALL,1)
REAL(KIND=JPRB)            :: D2WTH_fil(NSEQALL,1), D2WTH_inc(NSEQALL,1)

! SAVE for OpenMP
INTEGER(KIND=JPIM),SAVE    :: ISEQ, I
REAL(KIND=JPRD),SAVE       :: PSTOALL
REAL(KIND=JPRB),SAVE       :: DSTOALL, DSTO_add, DWTH_add
!================================================
P0GLBRIVSTO=0._JPRD
P0GLBFLDSTO=0._JPRD
P0GLBFLDARE=0._JPRD
P0GLBSTOPRE2=0._JPRD
P0GLBSTONEW2=0._JPRD

! [1] Assume all waters in river channel
!$OMP PARALLEL DO SIMD PRIVATE(DSTOALL,PSTOALL) REDUCTION(+:P0GLBSTOPRE2)
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBSTOPRE2     = P0GLBSTOPRE2 + PSTOALL

  D2FLDFRC(ISEQ,1) = 0._JPRB
  D2FLDDPH(ISEQ,1) = 0._JPRB

  D2STO_fil(ISEQ,1) = D2RIVSTOMAX(ISEQ,1)
  D2WTH_fil(ISEQ,1) = D2RIVWTH(ISEQ,1)
  D2WTH_inc(ISEQ,1) = D2GRAREA(ISEQ,1) / D2RIVLEN(ISEQ,1) * DFRCINC
END DO
!$OMP END PARALLEL DO SIMD

! [2] Check floodplain level from I=1 to NLFP. Make I-NLFP loop outside for parallel computing (SIMD/Vector)
DO I=1, NLFP
!$OMP PARALLEL DO SIMD PRIVATE(DSTOALL,DSTO_add,DWTH_add,PSTOALL) 
  DO ISEQ=1, NSEQALL
    PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
    DSTOALL = REAL( PSTOALL, KIND=JPRB )

    DSTO_add = DSTOALL - D2STO_fil(ISEQ,1)
    DSTO_add = MAX( DSTO_add, 0._JPRB )
    DWTH_add = -D2WTH_fil(ISEQ,1) + &
&    ( D2WTH_fil(ISEQ,1)**2._JPRB + 2._JPRB * DSTO_add / D2RIVLEN(ISEQ,1) / D2FLDGRD(ISEQ,1,I) )**0.5_JPRB
    DWTH_add = MIN( DWTH_add, D2WTH_inc(ISEQ,1) )
    DWTH_add = MAX( DWTH_add, 0._JPRB )              !! modify v4.04

    D2FLDDPH(ISEQ,1) = D2FLDDPH(ISEQ,1) + DWTH_add * D2FLDGRD(ISEQ,1,I)   !! DWTH_add=0 when flddpth < layer fldhgt
    D2FLDFRC(ISEQ,1) = D2FLDFRC(ISEQ,1) + DWTH_add / D2WTH_inc(ISEQ,1) / NLFP

    !! Update downside floodplain step storage/depth/width 
    D2STO_fil(ISEQ,1) = D2FLDSTOMAX(ISEQ,1,I)
    D2WTH_fil(ISEQ,1) = D2WTH_fil(ISEQ,1) + D2WTH_inc(ISEQ,1)
  ENDDO
  !$OMP END PARALLEL DO SIMD
ENDDO

!! [3] flood extent saturated case
!$OMP PARALLEL DO SIMD PRIVATE(DSTOALL,DSTO_add,PSTOALL) 
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)
  DSTO_add = DSTOALL - D2STO_fil(ISEQ,1)
  DSTO_add = MAX( DSTO_add, 0._JPRB )
  D2FLDDPH(ISEQ,1) = D2FLDDPH(ISEQ,1) + DSTO_add / D2WTH_fil(ISEQ,1) / D2RIVLEN(ISEQ,1)  !! DSTO_add=0 when not saturated
ENDDO
!$OMP END PARALLEL DO SIMD

!! [4] Floodplain stage diagnose
!$OMP PARALLEL DO SIMD PRIVATE(PSTOALL) 
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)

  P2RIVSTO(ISEQ,1) = D2RIVSTOMAX(ISEQ,1) + D2RIVLEN(ISEQ,1)*D2RIVWTH(ISEQ,1)*D2FLDDPH(ISEQ,1)
  P2RIVSTO(ISEQ,1) = MIN( P2RIVSTO(ISEQ,1), PSTOALL ) 
  D2RIVDPH(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1),KIND=JPRB) / (D2RIVLEN(ISEQ,1)*D2RIVWTH(ISEQ,1)) 

  P2FLDSTO(ISEQ,1) = PSTOALL - P2RIVSTO(ISEQ,1)
  P2FLDSTO(ISEQ,1) = MAX( P2FLDSTO(ISEQ,1), 0._JPRD )

  D2FLDFRC(ISEQ,1) = MAX( D2FLDFRC(ISEQ,1),0._JPRB)
  D2FLDFRC(ISEQ,1) = MIN( D2FLDFRC(ISEQ,1),1._JPRB)
  D2FLDARE(ISEQ,1) = D2GRAREA(ISEQ,1)*D2FLDFRC(ISEQ,1)
ENDDO
!$OMP END PARALLEL DO SIMD

!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTONEW2,P0GLBRIVSTO,P0GLBFLDSTO,P0GLBFLDARE)
DO ISEQ=1, NSEQALL
  D2SFCELV(ISEQ,1)  = D2RIVELV(ISEQ,1) + D2RIVDPH(ISEQ,1)
  D2STORGE(ISEQ,1)  = REAL(P2RIVSTO(ISEQ,1)+P2FLDSTO(ISEQ,1),KIND=JPRB)
  P0GLBSTONEW2      = P0GLBSTONEW2+ P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBRIVSTO       = P0GLBRIVSTO + P2RIVSTO(ISEQ,1)
  P0GLBFLDSTO       = P0GLBFLDSTO + P2FLDSTO(ISEQ,1)
  P0GLBFLDARE       = P0GLBFLDARE + D2FLDARE(ISEQ,1)
ENDDO
!$OMP END PARALLEL DO SIMD

END SUBROUTINE CMF_OPT_FLDSTG_ES
!####################################################################
END MODULE CMF_CALC_FLDSTG_MOD
